# 导入PyTorch库，用于构建和训练神经网络
import torch
# 导入matplotlib.pyplot库，用于绘制损失变化曲线
import matplotlib.pyplot as plt

# 定义输入数据x，形状为(4,1)，表示4个样本，每个样本1个特征
x_data = torch.Tensor([[1.0], [2.0], [3.0], [4.0]])  # 注意这里是双层括号，确保维度正确
# 定义目标输出数据y，形状与x_data一致，用于监督学习
y_data = torch.Tensor([[4.0], [5.0], [6.0], [7.0]])

# 定义线性回归模型类，继承自PyTorch的Module类
class LinearModel(torch.nn.Module):
    # 构造函数，初始化模型参数
    def __init__(self):
        # 调用父类的构造函数
        super(LinearModel, self).__init__()
        # 定义线性层：输入特征数为1，输出特征数为1
        self.linear = torch.nn.Linear(1, 1)

    # 前向传播函数，定义模型的计算过程
    def forward(self, x):
        # 通过线性层计算预测值y_pred
        y_pred = self.linear(x)
        return y_pred

# 创建模型实例
model = LinearModel()

# 定义损失函数：均方误差损失，reduction='sum'表示对所有样本的误差求和
# 注意：原代码中的size_average=False已过时，建议改为reduction='sum'
criterion = torch.nn.MSELoss(size_average=False)
# 定义优化器：平均随机梯度下降(ASGD)，学习率为0.01，优化模型的所有参数
optimizer = torch.optim.ASGD(model.parameters(), lr=0.01)

# 创建列表存储每轮的损失值，用于后续可视化
loss_list = []
# 创建列表存储轮次，用于绘图的x轴
epoch_list = []

# 训练模型，共迭代1000轮
for epoch in range(1000):
    # 前向传播：使用当前模型参数计算预测值
    y_pred = model(x_data)
    # 计算损失：预测值与真实值之间的均方误差
    loss = criterion(y_pred, y_data)
    # 打印当前轮次和对应的损失值
    print("epoch:", epoch, "loss:", loss.item())
    # 将当前损失值添加到列表中
    loss_list.append(loss.item())
    # 将当前轮次添加到列表中
    epoch_list.append(epoch)

    # 清零梯度：防止梯度累积
    optimizer.zero_grad()
    # 反向传播：计算损失对各参数的梯度
    loss.backward()
    # 参数更新：根据梯度和学习率更新模型参数
    optimizer.step()

# 输出训练得到的权重w
print("w = ", model.linear.weight.item())
# 输出训练得到的偏置b
print("b = ", model.linear.bias.item())

# 定义测试输入x_test，值为4.0
x_test = torch.tensor([4.0])
# 使用训练好的模型进行预测
y_test = model(x_test)
# 输出预测结果
print("y_test:", y_test.item())

# 绘制损失值随训练轮次的变化曲线
plt.plot(epoch_list, loss_list)
# 设置x轴标签
plt.xlabel("epoch")
# 设置y轴标签
plt.ylabel("loss")
# 显示图像
plt.show()